import numpy as np
import os
import random
import itertools


class multimodal(object):

    def __init__(self, qValue, volFrac):
        self.qValue = qValue
        self.volFrac = volFrac

    def load(self, fileName):
        print('LOADING ORIGINAL SPPARKS SITE FILE......')
        self.featureArray = np.loadtxt('{}.spparks'.format(fileName), skiprows=9).astype(np.int64)[:, 1]

        f = open('{}.spparks'.format(fileName), 'r')
        for i, line in enumerate(f, 0):
            if i == 2:
                self.numSites = int(line.split()[0])
            elif i == 4:
                self.xDim = int(line.split()[1])
            elif i == 5:
                self.yDim = int(line.split()[1])
            elif i == 6:
                self.zDim = int(line.split()[1])
                break
        f.close()

        self.numPrecipitates = self.numSites * self.volFrac
        self.numPrecipitates = np.round(self.numPrecipitates).astype(np.int64)

    def rand(self, projectName):

        self.projectName = projectName

        print('-------------------------------------------------------')
        print('GENERATING RANDOM PRECIPITATES DISTRIBUTION......')
        pSiteIdx = random.sample(range(0, self.numSites), self.numPrecipitates)

        featureRandArray = np.copy(self.featureArray)
        featureRandArray[pSiteIdx] = self.qValue + 1
        featureRandArray = featureRandArray.reshape(-1, 1)
        valuesArray = np.hstack((np.arange(1, self.numSites + 1).reshape(-1, 1), featureRandArray))

        print('WRITTING SPPARKS SITE FILE......')
        f = open('{}.spparks'.format(projectName), 'w', newline='\n')
        print('-', file=f)
        print('3 dimension', file=f)
        print('{} sites'.format(self.numSites), file=f)
        print('26 max neighbors', file=f)
        print('0 {} xlo xhi'.format(self.xDim), file=f)
        print('0 {} ylo yhi'.format(self.yDim), file=f)
        print('0 {} zlo zhi'.format(self.zDim), file=f)
        print('', file=f)
        print('Values', file=f)
        print('', file=f)
        np.savetxt(f, valuesArray, fmt='%d', delimiter='\t')
        f.close()

    def quadrant(self, projectName):

        self.projectName = projectName

        numPrecipitatesQuad = int((self.numPrecipitates - (self.numPrecipitates % 5)) / 5)

        idxArray = np.arange(0, self.numSites).reshape((self.zDim, self.xDim, self.yDim))
        bottomIdx = idxArray[:self.zDim // 2, ...]
        topIdx = idxArray[self.zDim // 2:, ...]
        print('-------------------------------------------------------')
        print('GENERATING QUADRANT PRECIPITATES DISTRIBUTION......')
        firstQuaterIdx = np.split(bottomIdx, [self.xDim // 2], axis=2)[0].flatten()
        secondQuaterIdx = np.split(bottomIdx, [self.xDim // 2], axis=2)[1].flatten()
        thirdQuaterIdx = np.split(topIdx, [self.xDim // 2], axis=2)[0].flatten()
        fourthQuaterIdx = np.split(topIdx, [self.xDim // 2], axis=2)[1].flatten()

        featureQuadArray = np.copy(self.featureArray)
        pSiteIdx = firstQuaterIdx[random.sample(range(0, len(firstQuaterIdx)), numPrecipitatesQuad * 2)]
        featureQuadArray[pSiteIdx] = self.qValue + 1

        pSiteIdx = secondQuaterIdx[random.sample(range(0, len(secondQuaterIdx)), numPrecipitatesQuad)]
        featureQuadArray[pSiteIdx] = self.qValue + 1

        pSiteIdx = thirdQuaterIdx[random.sample(range(0, len(thirdQuaterIdx)), numPrecipitatesQuad)]
        featureQuadArray[pSiteIdx] = self.qValue + 1

        pSiteIdx = fourthQuaterIdx[random.sample(range(0, len(fourthQuaterIdx)), numPrecipitatesQuad)]
        featureQuadArray[pSiteIdx] = self.qValue + 1

        featureQuadArray = featureQuadArray.reshape(-1, 1)
        valuesArray = np.hstack((np.arange(1, self.numSites + 1).reshape(-1, 1), featureQuadArray))

        print('WRITTING SPPARKS SITE FILE......')
        f = open('{}.spparks'.format(projectName), 'w', newline='\n')
        print('-', file=f)
        print('3 dimension', file=f)
        print('{} sites'.format(self.numSites), file=f)
        print('26 max neighbors', file=f)
        print('0 {} xlo xhi'.format(self.xDim), file=f)
        print('0 {} ylo yhi'.format(self.yDim), file=f)
        print('0 {} zlo zhi'.format(self.zDim), file=f)
        print('', file=f)
        print('Values', file=f)
        print('', file=f)
        np.savetxt(f, valuesArray, fmt='%d', delimiter='\t')
        f.close()

    def uniform(self, projectName):

        print('-------------------------------------------------------')
        print('GENERATING UNIFORM PRECIPITATES DISTRIBUTION......')

        numPrecipitatesUni = np.round(np.cbrt(self.numPrecipitates)).astype(np.int64)

        # DEFINE PADDING
        x = self.xDim - 8
        y = self.yDim - 8
        z = self.zDim - 8

        xStep = np.round(x / numPrecipitatesUni).astype(np.int64)
        yStep = np.round(y / numPrecipitatesUni).astype(np.int64)
        zStep = np.round(z / numPrecipitatesUni).astype(np.int64)

        xIdx = np.arange(5, x + xStep, xStep)
        yIdx = np.arange(5, y + yStep, yStep)
        zIdx = np.arange(5, z + zStep, zStep)

        if len(xIdx) - numPrecipitatesUni > 0:
            d = len(xIdx) - numPrecipitatesUni
            xIdx = xIdx[:-d]

        if len(yIdx) - numPrecipitatesUni > 0:
            d = len(yIdx) - numPrecipitatesUni
            yIdx = yIdx[:-d]

        if len(zIdx) - numPrecipitatesUni > 0:
            d = len(zIdx) - numPrecipitatesUni
            zIdx = zIdx[:-d]

        idxArray = np.arange(0, self.numSites).reshape(self.zDim, self.xDim, self.yDim)
        pSiteIdx = np.empty(0).astype(np.int64)
        for i in itertools.product(zIdx, yIdx, xIdx):
            pSiteIdx = np.append(pSiteIdx, idxArray[i[0], i[1], i[2]])

        featureUniArray = np.copy(self.featureArray)
        featureUniArray[pSiteIdx] = self.qValue + 1
        featureUniArray = featureUniArray.reshape(-1, 1)
        valuesArray = np.hstack((np.arange(1, self.numSites + 1).reshape(-1, 1), featureUniArray))

        print('WRITTING SPPARKS SITE FILE......')
        f = open('{}.spparks'.format(projectName), 'w', newline='\n')
        print('-', file=f)
        print('3 dimension', file=f)
        print('{} sites'.format(self.numSites), file=f)
        print('26 max neighbors', file=f)
        print('0 {} xlo xhi'.format(self.xDim), file=f)
        print('0 {} ylo yhi'.format(self.yDim), file=f)
        print('0 {} zlo zhi'.format(self.zDim), file=f)
        print('', file=f)
        print('Values', file=f)
        print('', file=f)
        np.savetxt(f, valuesArray, fmt='%d', delimiter='\t')
        f.close()


    def banded(self, projectName):

        print('-------------------------------------------------------')
        print('GENERATING BANDED PRECIPITATES DISTRIBUTION......')

        numPrecipitatesBanded = np.round(self.numPrecipitates / 11).astype(np.int64)

        z = self.zDim - 8
        zStep = np.round(z / 10).astype(np.int64)
        tempZIdx = np.arange(5, z + zStep, zStep)
        zIdx = np.empty(0).astype(np.int64)
        for i in tempZIdx:
            zIdx = np.append(zIdx, [i-1, i, i+1])

        zIdx = zIdx.reshape(-1, 3)

        idxArray = np.arange(0, self.numSites).reshape((self.zDim, self.xDim, self.yDim))
        pSiteIdx = np.empty(0).astype(np.int64)
        for i in range(zIdx.shape[0]):
            tempIdx = random.sample(range(0, self.xDim * self.yDim * 3), numPrecipitatesBanded)
            pSiteIdx = np.append(pSiteIdx, idxArray[zIdx[i], ...].flatten()[tempIdx])

        featureBandedArray = np.copy(self.featureArray)
        featureBandedArray[pSiteIdx] = self.qValue + 1
        featureBandedArray = featureBandedArray.reshape(-1, 1)
        valuesArray = np.hstack((np.arange(1, self.numSites + 1).reshape(-1, 1), featureBandedArray))

        print('WRITTING SPPARKS SITE FILE......')
        f = open('{}.spparks'.format(projectName), 'w', newline='\n')
        print('-', file=f)
        print('3 dimension', file=f)
        print('{} sites'.format(self.numSites), file=f)
        print('26 max neighbors', file=f)
        print('0 {} xlo xhi'.format(self.xDim), file=f)
        print('0 {} ylo yhi'.format(self.yDim), file=f)
        print('0 {} zlo zhi'.format(self.zDim), file=f)
        print('', file=f)
        print('Values', file=f)
        print('', file=f)
        np.savetxt(f, valuesArray, fmt='%d', delimiter='\t')
        f.close()

    def changeVolFrac(self, newVolFrac):
        self.volFrac = newVolFrac
        self.numPrecipitates = self.numSites * self.volFrac
        self.numPrecipitates = np.round(self.numPrecipitates).astype(np.int64)

    def writesSPKINP(self, withP=True):
        print('WRITTING SPPARKS INPUT FILE......')
        f = open('{}.in'.format(self.projectName), 'w', newline='\n')
        print('log 			spk-${N}.log', file=f)
        print('seed            56789', file=f)
        if withP:
            print('app_style       potts/pin {}'.format(int(self.qValue)), file=f)
        else:
            print('app_style       potts/neighonly {}'.format(int(self.qValue)), file=f)
        print('dimension       3', file=f)
        print('lattice		    sc/26n 1.0', file=f)
        print('region          square block 0 {} 0 {} 0 {}'.format(self.xDim, self.yDim, self.zDim), file=f)
        print('create_box	    square', file=f)
        print('create_sites    box', file=f)
        print('read_sites      spk-${N}.spparks', file=f)
        print('temperature	    0.85', file=f)
        print('solve_style     tree', file=f)
        print('sector          yes', file=f)
        print('diag_style      energy', file=f)
        print('diag_style      propensity', file=f)
        print('diag_style      cluster stats no delt 50.0 filename spk-${N}.cluster', file=f)
        print('stats           1 loglinfreq 10 10.0', file=f)
        print('dump            1 text 1.0 spk-${N}.dump.*', file=f)
        print('dump_modify     1 loglinfreq 10 10.0 flush yes', file=f)
        print('run             200000.0', file=f)
        f.close()

    def writesBashFile(self):
        f = open('{}.sh'.format(self.projectName), 'w', newline='\n')
        print('#!/bin/bash', file=f)
        print('let N={}'.format(self.projectName[4:]), file=f)
        print('mpirun -np 16 spk -var N $N < {}.in'.format(self.projectName), file=f)
        f.close()
